'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2020-04-07 22:02:52
'''
# -*- coding: UTF-8 -*-
# nohup python -u galaxy_batch.py 8234 >> result &
from subprocess import *
import random, time, json, sys
# from kafka import KafkaProducer
# from kafka import KafkaConsumer
# from kafka import TopicPartition
# import sys
# sys.path.append('../util')
# from feature_handler import get_fea_index, fea_code


class BatchTrain():
    def __init__(self, model_name, batch_size, shuffle_size, watch_num, neg_rate, ps_url, process_num, dates, f_path, run_shell, net_conf, labels, random_sleep=0):
        self.model_name = model_name
        self.batch_size = batch_size
        self.process_num = process_num
        self.shuffle_size = shuffle_size
        self.watch_num = watch_num
        self.neg_rate = neg_rate
        self.ps_url = ps_url
        self.run_shell = run_shell
        self.dates = dates
        self.f_path = f_path
        self.random_sleep = random_sleep
        self.net_conf = net_conf
        self.labels = labels
        self.shell_p = self.run_shell + " -layer /data1/xuejian/sync/model_zoo/models/galaxy/" + self.net_conf + " -mn " + self.model_name + " -ps_url " + self.ps_url + ':' + str(sys.argv[1]) + " -w " + str(self.watch_num) + " -s " + str(self.shuffle_size) + " -b " + str(self.batch_size) + " -ne " + str(self.neg_rate) + " -labels " + self.labels
        print('shell = ', self.shell_p)
        
        


    def train(self):
        for date in self.dates:
            print(date)
            proc_dict = {}
            num_list = [0] * self.process_num
            print("num_list = ", num_list) 
            for i in range(self.process_num):
                if self.process_num <= 1:
                    f_date = self.f_path + date
                else:
                    f_date = self.f_path + date + '/' + str(i)
                if i != 0:
                    time.sleep(random.randint(0, self.random_sleep))
                proc = Popen('cat ' + f_date + ' | ' + self.shell_p, shell=True)
                print('process start')
                proc_dict[i] = proc
            for _, p in proc_dict.items():
                p.wait()

    def train2(self):
        
        proc = Popen(self.shell_p, shell=True, stdin=PIPE)
        print('process start')
        for date in self.dates:
            print(date)
            with open(self.f_path + date, 'rb') as f_date:
                print(self.f_path + date)
                line = f_date.readline()
                while line != b'':
                    if line == b'':
                        break
                    proc.stdin.write(line)
                    proc.stdin.flush()
                    line = f_date.readline()
        
        proc.stdin.close()
        proc.wait()



if __name__ == '__main__':
    # for the daily update argv1 is port, 2 is date and 3 is model_name
    dates = []
    dates.append(sys.argv[2])

    cut_path = '/data3/training_data/galaxy_data_cut/'
    single_path = '/data3/training_data/galaxy_data_new/'
    nums = 6
    if nums <= 1:
        real_path = single_path
        watch_num = 200000
    else:
        real_path = cut_path
        watch_num = 50000
    # single_path = '/data1/zangruozhou/sync/transfer/'
    # ps_url = '172.21.39.3' #offlineps-03, offline train daily update  port:8234
    # ps_url = '172.21.39.3' #offlineps-03, offline train  port:1111 the real time mode
    ps_url = '172.21.39.3' #offlineps-03, offline train  port:3333 the tmp mode, port 5555 for the deepfm_id, daily update.
    # ps_url = ''
    # net_name = 'success_deepfm.net'
    model_name = sys.argv[3]
    net_name = model_name + '.net'
    # net_name = 'success_deepfm_id.net' #for the deepfm_id model.
    # model_name = 'success_fm'
    batch_train = BatchTrain(model_name, 200, 10000, watch_num, 0.99, ps_url, nums, dates, real_path, '/data1/xuejian/sync/galaxy/galaxy_old', net_name, 'one', 15)
    if nums <= 1:
        batch_train.train2()
    else:
        batch_train.train()